package org.test.boot.rocketmq.producer;

import org.apache.rocketmq.client.producer.SendCallback;
import org.apache.rocketmq.client.producer.SendResult;
import org.apache.rocketmq.spring.annotation.RocketMQTransactionListener;
import org.apache.rocketmq.spring.core.RocketMQLocalTransactionListener;
import org.apache.rocketmq.spring.core.RocketMQLocalTransactionState;
import org.apache.rocketmq.spring.core.RocketMQTemplate;
import org.apache.rocketmq.spring.support.RocketMQHeaders;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.boot.CommandLineRunner;
import org.springframework.messaging.Message;
import org.springframework.messaging.MessageHeaders;
import org.springframework.messaging.support.MessageBuilder;
import org.springframework.stereotype.Component;
import org.springframework.util.MimeTypeUtils;
import org.test.boot.rocketmq.config.MqProperties;
import org.test.boot.rocketmq.domain.OrderPaidEvent;
import org.test.boot.rocketmq.domain.User;

import javax.annotation.Resource;
import java.math.BigDecimal;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicInteger;

@Component
public class ProductProducer implements CommandLineRunner {
    private static final Logger LOGGER = LoggerFactory.getLogger(ProductProducer.class);
    @Resource
    private RocketMQTemplate rocketMQTemplate;

    @Resource
    private MqProperties mqProperties;

    @Resource(name = "extRocketMQTemplate")
    private RocketMQTemplate extRocketMQTemplate;

    @Override
    public void run(String... args) throws Exception {
        String topic = mqProperties.getTopic();
        SendResult sendResult = rocketMQTemplate.syncSend(topic, "Hello, World!");
        LOGGER.info("syncSend1 to topic {} sendResult={} %n", topic, sendResult);

        String userTopic = mqProperties.getUserTopic();
        // 发送对象消息
        sendResult = rocketMQTemplate.syncSend(userTopic, new User().setUserAge((byte) 18).setUserName("Zhang"));
        LOGGER.info("syncSend1 to topic {} sendResult={} %n", userTopic, sendResult);

        sendResult = rocketMQTemplate.syncSend(userTopic, MessageBuilder.withPayload(
                new User().setUserAge((byte) 21).setUserName("Lester")).setHeader(MessageHeaders.CONTENT_TYPE, MimeTypeUtils.APPLICATION_JSON_VALUE).build());
        LOGGER.info("syncSend1 to topic {} sendResult={} %n", userTopic, sendResult);

        // Use the extRocketMQTemplate
        sendResult = extRocketMQTemplate.syncSend(topic, MessageBuilder.withPayload("Hello, World!2222".getBytes()).build());
        System.out.printf("extRocketMQTemplate.syncSend1 to topic %s sendResult=%s %n", topic, sendResult);

        // Send user-defined object
        rocketMQTemplate.asyncSend(mqProperties.getOrderPaidTopic(), new OrderPaidEvent("T_001", new BigDecimal("88.00")), new SendCallback() {
            @Override
            public void onSuccess(SendResult var1) {
                LOGGER.info("async onSuccess SendResult={}", var1);
            }

            @Override
            public void onException(Throwable var1) {
                LOGGER.info("async onException Throwable", var1);
            }

        });

        // Send message with special tag
        // tag0 will not be consumer-selected
        rocketMQTemplate.convertAndSend(topic + ":tag0", "I'm from tag0");
        LOGGER.info("syncSend topic {} tag {}", topic, "tag0");
        rocketMQTemplate.convertAndSend(topic + ":tag1", "I'm from tag1");
        LOGGER.info("syncSend topic {} tag {}", topic, "tag1");

        // Send transactional messages using rocketMQTemplate
        testRocketMQTemplateTransaction();
    }

    private void testRocketMQTemplateTransaction() {
        String[] tags = new String[]{"TagA", "TagB", "TagC", "TagD", "TagE"};
        for (int i = 0; i < 10; i++) {
            try {
                Message<String> msg = MessageBuilder.withPayload("rocketMQTemplate transactional message " + i).
                        setHeader(RocketMQHeaders.TRANSACTION_ID, "KEY_" + i).build();
                SendResult sendResult = rocketMQTemplate.sendMessageInTransaction(
                        mqProperties.getTopic() + ":" + tags[i % tags.length], msg, null);
                LOGGER.info("------rocketMQTemplate send Transactional msg body = {} , sendResult={}",
                        msg.getPayload(), sendResult.getSendStatus());
                Thread.sleep(10);
            } catch (Exception e) {
                LOGGER.error("send transactional message exception", e);
            }
        }
    }

    @RocketMQTransactionListener
    static class TransactionListenerImpl implements RocketMQLocalTransactionListener {
        private final AtomicInteger transactionIndex = new AtomicInteger(0);

        private final ConcurrentHashMap<String, Integer> localTrans = new ConcurrentHashMap<>();

        @Override
        public RocketMQLocalTransactionState executeLocalTransaction(Message msg, Object arg) {
            String transId = (String) msg.getHeaders().get(RocketMQHeaders.TRANSACTION_ID);
            LOGGER.info("#### executeLocalTransaction is executed, msgTransactionId={}", transId);
            int value = transactionIndex.getAndIncrement();
            int status = value % 3;
            localTrans.put(transId, status);
            if (status == 0) {
                // Return local transaction with success(commit), in this case,
                // this message will not be checked in checkLocalTransaction()
                LOGGER.info("    # COMMIT # Simulating msg {} related local transaction exec succeeded! ###", msg.getPayload());
                return RocketMQLocalTransactionState.COMMIT;
            } else if (status == 1) {
                // Return local transaction with failure(rollback) , in this case,
                // this message will not be checked in checkLocalTransaction()
                LOGGER.info("    # ROLLBACK # Simulating {} related local transaction exec failed!", msg.getPayload());
                return RocketMQLocalTransactionState.ROLLBACK;
            }

            LOGGER.info("    # UNKNOW # Simulating {} related local transaction exec UNKNOWN!", msg.getPayload());
            return RocketMQLocalTransactionState.UNKNOWN;
        }

        @Override
        public RocketMQLocalTransactionState checkLocalTransaction(Message msg) {
            String transId = (String) msg.getHeaders().get(RocketMQHeaders.TRANSACTION_ID);
            RocketMQLocalTransactionState retState = RocketMQLocalTransactionState.COMMIT;
            Integer status = localTrans.get(transId);
            if (null != status) {
                switch (status) {
                    case 0:
                        retState = RocketMQLocalTransactionState.COMMIT;
                        break;
                    case 1:
                        retState = RocketMQLocalTransactionState.ROLLBACK;
                        break;
                    case 2:
                        retState = RocketMQLocalTransactionState.UNKNOWN;
                        break;
                }
            }
            LOGGER.info("------ !!! checkLocalTransaction is executed once," +
                            " msgTransactionId={}, TransactionState={} status={} %n",
                    transId, retState, status);
            return retState;
        }
    }
}
